﻿// ==++==
// 
//   Copyright (c) Microsoft Corporation.  All rights reserved.
// 
// ==--==
using System;
using System.Collections.Generic;
using System.Text;
using Microsoft.Pex.Engine.ComponentModel;
using Microsoft.ExtendedReflection.Interpretation.Interpreter;
using Microsoft.ExtendedReflection.Metadata;
using Microsoft.ExtendedReflection.Interpretation;
using Microsoft.ExtendedReflection.Reasoning.ExecutionNodes;
using Microsoft.ExtendedReflection.Collections;
using Microsoft.ExtendedReflection.Utilities;
using Microsoft.ExtendedReflection.ComponentModel;
using Microsoft.ExtendedReflection.Utilities.Safe.Diagnostics;
using Microsoft.Pex.Fitness.Interfaces;
using Microsoft.ExtendedReflection.Interpretation.Interpreter.Monitoring;
using Microsoft.Pex.Engine.ExecutionNodes;

namespace Microsoft.Pex.Fitness.Implementation
{
    internal class PexFitnessPathManager
        : PexPathComponentBase
        , IPexFitnessPathManager
    {
        const int FitnessValueForSymbolicNodeWithUncoveredBranchSide = 0;

        IPexFitnessExplorationManager fitnessExplorationManager;
        TermManager termManager;
        IInstructionInterpreterEventSource instructionInterpreterEventSource;
        IPexCurrentExecutionNodeProvider currentExecutionNodeProvider;
        IExecutionNode flippedNode;
        SafeSet<int> targetsImprovedInThisPath = new SafeSet<int>();
        SafeIntDictionary<int> theOtherPredicateSideIds = new SafeIntDictionary<int>();

        //If we set the flag below as false, then we will remove all the fitness info for a target associated with nodes 
        //from the memory once the target is covered. Then when we output the graph, we don't get the
        //historical fitness info for nodes that has been removed.
        //So here we set it as true so we don't remove info after we cover a target
        static bool keepAlreadyCoveredTargetOrBranchFitnessForDebug =
#if DEBUG
            true;
#else
            false;
#endif

        protected override void Initialize()
        {
            base.Initialize();
            this.fitnessExplorationManager = ServiceProviderHelper.GetService<IPexFitnessExplorationManager>(this.Site);
            this.termManager = this.ExplorationServices.TermManager;
            this.instructionInterpreterEventSource = this.ExplorationServices.InstructionInterpreterEventSource;
            SafeDebug.AssumeNotNull(this.instructionInterpreterEventSource, "this.instructionInterpreterEventSource");
            this.instructionInterpreterEventSource.BinaryOperationHandler += new EventHandler<InstructionInterpreterBinaryOperationEventArgs>(InstructionInterpreterEventSource_BinaryOperationHandler);
            this.currentExecutionNodeProvider = this.PathServices.CurrentExecutionNodeProvider;
            this.flippedNode = this.PathServices.InputModelProvider.InputModel.ParentOfTarget;
        }

        public override void Dispose()
        {
            var instructionInterpreterEventSource = this.instructionInterpreterEventSource;
            if (instructionInterpreterEventSource != null)
            {
                this.instructionInterpreterEventSource = null;
                instructionInterpreterEventSource.BinaryOperationHandler -= new EventHandler<InstructionInterpreterBinaryOperationEventArgs>(InstructionInterpreterEventSource_BinaryOperationHandler);
            }
            base.Dispose();
        }

        public void OnPathFinished()
        {
            var reversedNodes = this.currentExecutionNodeProvider.ReversedNodes;
            TermManager termManager = this.termManager;
            IExecutionNode subsequentNodeInPath = null;
            foreach (IExecutionNode node in reversedNodes)//walking backward in path
            {
                if (subsequentNodeInPath != null)
                {
                    Term condition;
                    foreach (var kvp in node.Successors)//walk through each node in the current execution path
                        if (kvp.Value == subsequentNodeInPath)
                        {
                            condition = kvp.Key;
                            //the 'condition' is how we got from 'node' to 'lastNode'
                            bool negated = false;
                            Term innerNegatedValue;
                            if (termManager.TryGetInnerLogicallyNegatedValue(condition, out innerNegatedValue))
                            {
                                negated = true;
                                condition = innerNegatedValue;
                            }

                            BinaryOperator @op;
                            Term left, right;
                            if (termManager.TryGetBinary(condition, out @op, out left, out right))
                            {
                                //we hashCode with 48572305 to avoid collision with the programCounter used in binary predicate tracking
                                this.CalculateAndSetFitnessValueForBranch(
                                    node.CallStackTrace,
                                    node.CodeLocation.Method,
                                    HashCodeHelper.Combine(node.CodeLocation.Offset, 48572305),
                                    op, left, right, node,
                                    negated ^ IsLeft(node.CodeLocation, subsequentNodeInPath.InCodeBranch));
                            }
                            break;
                        }
                }
            }

#if DEBUG
            var currentPathNumber = this.fitnessExplorationManager.CurrentPathNumber;
            this.fitnessExplorationManager.NextPathNumber();
            if (this.fitnessExplorationManager.IsLoggingEnabled)
            {
                this.Log.LogMessage(
                            "fitness",
                            "Finished {0}th path ================================", currentPathNumber);
            }
#endif
            subsequentNodeInPath = this.currentExecutionNodeProvider.Current;
            //this.Log.LogMessage(
            //    "fitness",
            //    "last node: {0}",
            //    subsequentNodeInPath);

            foreach (var kvp in this.theOtherPredicateSideIds)
            {
                this.fitnessExplorationManager.SetFitness(
                    this.targetsImprovedInThisPath,
                    this.flippedNode,
                    reversedNodes,
                    kvp.Key,    //our target is the other side of the taken predicate side
                    kvp.Value);
            }
        }

        //Since we don't have a precise way to decide whether a particular branch is left branch or right branch,
        //we use the following mechnism to distinguish the two branches of one branching point:
        //we consider the first encountered branch as the left branch during exploration.
        //Although it may not be the real left branch, at least we can distinguish the left and right branches, which
        //is good enough for our context.
        static readonly SafeEquatableDictionary<CodeLocation, CodeBranch> leftBranches = new SafeEquatableDictionary<CodeLocation, CodeBranch>();
        static bool IsLeft(CodeLocation location, CodeBranch branch)
        {
            CodeBranch leftBranch;
            if (leftBranches.TryGetValue(location, out leftBranch))
                return leftBranch == branch;
            else
            {
                leftBranches[location] = branch;
                return true;
            }
        }

        private void CalculateAndSetFitnessValueForBranch(
            ICallStackTrace callStackTrace,
            MethodDefinition method, int programCounter,
            BinaryOperator op, Term left, Term right,
            IExecutionNode currentNode,
            bool negated)
        {
            //we try to give priority to a node 
            //whose the other side is not covered to flip
            //note that we don't need to check whether the node involves symbolic values
            //since it should; otherwise it won't be considered as a node

            TermManager termManager = this.termManager;

            Term leftConcreteValue, rightConcreteValue;
            if (termManager.IsValue(left))
                leftConcreteValue = left;
            else
                leftConcreteValue = this.PathServices.InputModelProvider.InputModel.Model.GetValue(left);

            if (termManager.IsValue(right))
                rightConcreteValue = right;
            else
                rightConcreteValue = this.PathServices.InputModelProvider.InputModel.Model.GetValue(right);

            Term result = termManager.Binary(op, leftConcreteValue, rightConcreteValue);
            bool taken = termManager.IsNotDefaultValue(result);//taken the positive side of the perdicate
            if (negated)
                taken = !taken;

            int takenPredicateSideId, theOtherPredicateSideId;
            takenPredicateSideId = this.fitnessExplorationManager.GetTargetHash(callStackTrace, method, programCounter, taken);
            theOtherPredicateSideId = this.fitnessExplorationManager.GetTargetHash(callStackTrace, method, programCounter, !taken);

#if DEBUG
            if (!this.fitnessExplorationManager.TargetIDMap.ContainsKey(takenPredicateSideId))
                this.fitnessExplorationManager.TargetIDMap.Add(takenPredicateSideId, " Branch " + taken + "@" + method.ShortName + programCounter + "\n           " + left + " " + op + " " + right);

            if (!this.fitnessExplorationManager.TargetIDMap.ContainsKey(theOtherPredicateSideId))
                this.fitnessExplorationManager.TargetIDMap.Add(theOtherPredicateSideId, " Branch " + !taken + "@" + method.ShortName + programCounter + "\n           " + left + " " + op + " " + right);
#endif

            //this.Log.LogMessage("fitness", "Processing branch (Counter: {7}) {0} :({1} {2} {3}) SideID: {4}/{5} after last node {6}",
            //    taken, leftConcreteValue, op, rightConcreteValue, takenPredicateSideId, theOtherPredicateSideId, this.PathServices.CurrentExecutionNodeProvider.Current, programCounter.ToString("x8"));

            if (!this.fitnessExplorationManager.HasCoveredTarget(takenPredicateSideId))
            {
                this.fitnessExplorationManager.AddCoveredTarget(takenPredicateSideId);
                this.fitnessExplorationManager.SetCoveredTarget(this.flippedNode, takenPredicateSideId);

                if (!keepAlreadyCoveredTargetOrBranchFitnessForDebug)
                {
                    //we next remove all fitness values and fitness gain for the takenPredicateSideId
                    this.fitnessExplorationManager.RemoveFitnessAfterTargetCoverage(takenPredicateSideId);
                    this.fitnessExplorationManager.RemoveFitnessGainAfterTargetCoverage(takenPredicateSideId);
                }

                this.fitnessExplorationManager.SendEventOnTargetCoverage(this.flippedNode, takenPredicateSideId);

                if (this.fitnessExplorationManager.IsLoggingEnabled)
                    this.Log.LogMessage("fitness", "BRANCH: Finish covering branch (Counter: {9}) side {0}:{1} ({2}): ({3} {4} {5}) SideID: {6}/{7} after last node {8}", method.ToString(), programCounter, taken, leftConcreteValue, op, rightConcreteValue, takenPredicateSideId, theOtherPredicateSideId, this.currentExecutionNodeProvider.Current, programCounter.ToString("x8"));
            }

            int fitnessValue = FitnessValueForSymbolicNodeWithUncoveredBranchSide;//we give the fitness value as 0 to give priority
            if (this.fitnessExplorationManager.HasCoveredTarget(theOtherPredicateSideId))
                //we have covered the other side, we don't need to compute fitness value
                return;
            //else
            //    this.Log.LogMessage("fitness", "BRANCH: Focus on {8} taken branch side (Counter: {9}) {0}:{1}: {2} {3} {4} SideID: {5}/{6} Set:{7}", method.ToString(), programCounter, left, op, right, takenPredicateSideId, theOtherPredicateSideId, fem.CoveredTargets.Count, !taken, programCounter.ToString("x8"));

            if (this.fitnessExplorationManager.IsLoggingEnabled)
                this.Log.LogMessage("fitness", "BRANCH: @@@@@@@@op={0}, left={1}, right={2}, fitness={3} @ {4}:{5} after last node {6}", op, leftConcreteValue, rightConcreteValue, fitnessValue, method.ToString(), programCounter, this.currentExecutionNodeProvider.Current);
            this.fitnessExplorationManager.SetFitnessValueToNode(this.targetsImprovedInThisPath, this.flippedNode, currentNode, theOtherPredicateSideId, fitnessValue, 0, 0);//we assign 0 to the last two args for now
        }

        void InstructionInterpreterEventSource_BinaryOperationHandler(object sender, InstructionInterpreterBinaryOperationEventArgs e)
        {
            var currentFrame = e.Interpreter.CurrentFrame;
            ICallStackTrace callStackTrace = e.Interpreter.CallStackTrace;
            Method method = currentFrame.Method;
            int programCounter = currentFrame.ProgramCounter;
            BinaryOperator op = e.Operator;
            Term left = e.Left;
            Term right = e.Right;

            SafeDebug.AssumeNotNull(method, "method");
            SafeDebug.AssumeNotNull(left, "left");
            SafeDebug.AssumeNotNull(right, "right");

            IExecutionNode currentNode = this.currentExecutionNodeProvider.Current;
            this.CalculateAndSetFitnessValueForPredicate(callStackTrace, method.Definition, programCounter, op, left, right, currentNode, false);
        }

        private bool IsInteger(Term term)
        {
            SafeDebug.AssumeNotNull(term, "term");
            switch (this.termManager.GetLayoutKind(term))
            {
                case LayoutKind.I1:
                case LayoutKind.I2:
                case LayoutKind.I4:
                case LayoutKind.I8:
                    return true;
                default:
                    return false;
            }
        }

        private void CalculateAndSetFitnessValueForPredicate(
            ICallStackTrace callStackTrace,
            MethodDefinition method, int programCounter,
            BinaryOperator op, Term left, Term right,
            IExecutionNode currentNode,
            bool negated)
        {
            if (this.IsInteger(left) &&
                this.IsInteger(right)
                &&
                (op == BinaryOperator.Ceq ||
                 op == BinaryOperator.Clt ||
                 op == BinaryOperator.Clt_Un)
               )
            {
                TermManager termManager = this.ExplorationServices.TermManager;
                //below is used to identify and filter out those predicates related to boolean comparison
                //they are not useful for us to keep track or target at since it doesn't have nice
                //fitness function: hit it or not hit
                if (op == BinaryOperator.Ceq)
                {
                    if (termManager.IsProperStackBoolean(left) &&
                        termManager.IsProperStackBoolean(right)) //left and right operands are boolean 
                    {
                        if (termManager.IsValue(left) && termManager.IsValue(right))//left and right operands are concrete
                        {
                            //note that for concrete non-boolean int values, when it is value of 1 or 0
                            //the above boolean checks still pass, we want to make sure that this condition
                            //check hasn't been observed to be non-boolean (not 0 or 1) concrete integer values
                            if (!this.fitnessExplorationManager.IsReachedNonBooleanComparison(method, programCounter))
                            {
                                //this.Log.LogMessage("fitness",
                                //    "Encounter a concrete boolean comparison: {0}:{1}", method.ToString(), programCounter);
                                return;
                            }
                        }
                        else
                        {
                            //this.Log.LogMessage("fitness",
                            //        "Encounter a symbolic boolean comparison: {0}:{1}", method.ToString(), programCounter);
                            return;
                        }
                    }
                    else
                    {
                        //left and rigth operands are not boolean, we remember that so that in the future
                        //when the operands have values of 0 or 1, we don't mistaken them as boolean comparison
                        this.fitnessExplorationManager.ReachedNonBooleanComparison(method, programCounter);
                    }
                }

                Term leftConcreteValue, rightConcreteValue;
                if (termManager.IsValue(left))
                    leftConcreteValue = left;
                else
                    leftConcreteValue = this.PathServices.InputModelProvider.InputModel.Model.GetValue(left);

                if (termManager.IsValue(right))
                    rightConcreteValue = right;
                else
                    rightConcreteValue = this.PathServices.InputModelProvider.InputModel.Model.GetValue(right);

                Term result = termManager.Binary(op, leftConcreteValue, rightConcreteValue);
                bool taken = termManager.IsNotDefaultValue(result);//taken the positive side of the perdicate
                if (negated)
                    taken = !taken;

                int takenPredicateSideId, theOtherPredicateSideId;
                takenPredicateSideId = this.fitnessExplorationManager.GetTargetHash(callStackTrace, method, programCounter, taken);
                theOtherPredicateSideId = this.fitnessExplorationManager.GetTargetHash(callStackTrace, method, programCounter, !taken);

#if DEBUG
                if (!this.fitnessExplorationManager.TargetIDMap.ContainsKey(takenPredicateSideId))
                    this.fitnessExplorationManager.TargetIDMap.Add(takenPredicateSideId, " Predicate " + taken + "@" + method.ShortName + programCounter + "\n           " + left + " " + op + " " + right);

                if (!this.fitnessExplorationManager.TargetIDMap.ContainsKey(theOtherPredicateSideId))
                    this.fitnessExplorationManager.TargetIDMap.Add(theOtherPredicateSideId, " Predicate " + !taken + "@" + method.ShortName + programCounter + "\n           " + left + " " + op + " " + right);
#endif
                //this.Log.LogMessage("fitness", "Processing predicate (Counter: {7}) {0} :({1} {2} {3}) SideID: {4}/{5} after last node {6}",
                //    taken, leftConcreteValue, op, rightConcreteValue, takenPredicateSideId, theOtherPredicateSideId, this.PathServices.CurrentExecutionNodeProvider.Current, programCounter.ToString("x8"));

                if (!this.fitnessExplorationManager.HasCoveredTarget(takenPredicateSideId))
                {
                    this.fitnessExplorationManager.AddCoveredTarget(takenPredicateSideId);
                    this.fitnessExplorationManager.SetCoveredTarget(this.flippedNode, takenPredicateSideId);
                    //we next remove all fitness values and fitness gain for the takenPredicateSideId
                    if (!keepAlreadyCoveredTargetOrBranchFitnessForDebug)
                    {
                        this.fitnessExplorationManager.RemoveFitnessAfterTargetCoverage(takenPredicateSideId);
                        this.fitnessExplorationManager.RemoveFitnessGainAfterTargetCoverage(takenPredicateSideId);
                    }

                    this.fitnessExplorationManager.SendEventOnTargetCoverage(this.flippedNode, takenPredicateSideId);

                    if (this.fitnessExplorationManager.IsLoggingEnabled)
                        this.Log.LogMessage("fitness", "PREDICATE: Finish covering (Counter: {9}) side {0}:{1} ({2}): ({3} {4} {5}) SideID: {6}/{7} after last node {8}", method.ToString(), programCounter, taken, leftConcreteValue, op, rightConcreteValue, takenPredicateSideId, theOtherPredicateSideId, currentNode, programCounter.ToString("x8"));
                }

                if (this.fitnessExplorationManager.HasCoveredTarget(theOtherPredicateSideId))
                    //we have covered the other side, we don't need to compute fitness value
                    return;
                //else
                //    this.Log.LogMessage("fitness", "PREDICATE: Focus on {8} taken side (Counter: {9}) {0}:{1}: {2} {3} {4} SideID: {5}/{6} Set:{7}", method.ToString(), programCounter, left, op, right, takenPredicateSideId, theOtherPredicateSideId, fem.CoveredTargets.Count, !taken, programCounter.ToString("x8"));

                switch (op)
                {
                    case BinaryOperator.Ceq://== 
                    case BinaryOperator.Clt://< 
                    case BinaryOperator.Clt_Un://< unsigned
                        Term diffValue = termManager.Sub(leftConcreteValue, rightConcreteValue);
                        int fitnessValue;
                        if (!termManager.TryGetI4Constant(diffValue, out fitnessValue))
                            fitnessValue = int.MaxValue;
                        if (fitnessValue < 0)
                            fitnessValue = -fitnessValue;

                        if ((fitnessValue == 0) && (op == BinaryOperator.Ceq)) //this can happen when 0 != 0
                            return;//we simply don't consider it

                        if (op != BinaryOperator.Ceq)
                        {
                            //for x < y, fitness function is |x - y| + 1
                            //since if we use x - y, when the fitness value is 0, the x < y is still not satisfied
                            fitnessValue++;
                        }

                        //int l, r;
                        //if (abs.TryGetI4Constant(left, out l) && //left is a concrete value
                        //    abs.TryGetI4Constant(right, out r)) //right is a concrete value

                        if (this.fitnessExplorationManager.IsLoggingEnabled)
                            this.Log.LogMessage("fitness", "PREDICATE @@@@@@@@op={0}, left={1}, right={2}, fitness={3} @ {4}:{5} after last node {6}", op, leftConcreteValue, rightConcreteValue, fitnessValue, method.ToString(), programCounter, this.currentExecutionNodeProvider.Current);

                        int otherFitnessValue;
                        if (!this.theOtherPredicateSideIds.TryGetValue(theOtherPredicateSideId, out otherFitnessValue))
                            otherFitnessValue = int.MaxValue;
                        if (fitnessValue < otherFitnessValue)
                            this.theOtherPredicateSideIds[theOtherPredicateSideId] = fitnessValue;
                        /*
                        this.fitnessExplorationManager.SetFitness(
                            this.targetsImprovedInThisPath,
                            this.flippedNode,
                            this.currentExecutionNodeProvider.ReversedNodes,
                            theOtherPredicateSideId,//our target is the other side of the taken predicate side
                            fitnessValue);
                        */
                        break;
                }
            }
        }
    }
}
